package mr

import (
	"log"
	"net"
	"net/http"
	"net/rpc"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Coordinator struct {
	// Your definitions here.
	MapTasks    map[int]*Task
	ReduceTasks map[int]*Task

	mutex sync.Mutex
}

// Your code here -- RPC handlers for the worker to call.
func (c *Coordinator) RequestTask(args *TaskArgs, reply *TaskReply) error {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	finishMap := true
	for _, t := range c.MapTasks {
		if !t.Assigned {
			t.Assigned = true
			reply.Task = *t
			break
		}

		if !t.Finished {
			finishMap = false
		}
	}

	if reply.Task.Assigned {
		go func(mapId int) {
			time.Sleep(10 * time.Second)

			c.mutex.Lock()
			defer c.mutex.Unlock()

			task := c.MapTasks[mapId]
			if !task.Finished {
				task.Assigned = false
			}
		}(reply.Task.MapId)

		return nil
	}

	if !finishMap {
		reply.Type = NO_TASK
		return nil
	}

	for _, t := range c.ReduceTasks {
		if !t.Assigned {
			t.Assigned = true
			reply.Task = *t
			break
		}
	}

	if reply.Task.Assigned {
		go func(reduceId int) {
			time.Sleep(10 * time.Second)

			c.mutex.Lock()
			defer c.mutex.Unlock()

			task := c.ReduceTasks[reduceId]
			if !task.Finished {
				task.Assigned = false
			}
		}(reply.Task.ReduceId)
	} else {
		reply.Type = NO_TASK
	}

	return nil
}

func (c *Coordinator) FinishTask(args *FinishArgs, reply *FinishReply) error {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	if args.Type == MAP_TASK {
		for _, filename := range args.OutputFilenames {
			tems := strings.Split(filename, "-")
			reduceId, err := strconv.Atoi(tems[2])
			if err != nil {
				log.Println(err)
			}

			task := c.ReduceTasks[reduceId]
			task.InputFilenames = append(task.InputFilenames, filename)
		}
		// log.Printf("%+v\n", args)
		task := c.MapTasks[args.MapId]
		task.Finished = true
	} else if args.Type == REDUCE_TASK {
		task := c.ReduceTasks[args.ReduceId]
		task.Finished = true
	}

	return nil
}

//
// an example RPC handler.
//
// the RPC argument and reply types are defined in rpc.go.
//
func (c *Coordinator) Example(args *ExampleArgs, reply *ExampleReply) error {
	reply.Y = args.X + 1
	return nil
}

//
// start a thread that listens for RPCs from worker.go
//
func (c *Coordinator) server() {
	rpc.Register(c)
	rpc.HandleHTTP()
	//l, e := net.Listen("tcp", ":1234")
	sockname := coordinatorSock()
	os.Remove(sockname)
	l, e := net.Listen("unix", sockname)
	if e != nil {
		log.Fatal("listen error:", e)
	}
	go http.Serve(l, nil)
}

//
// main/mrcoordinator.go calls Done() periodically to find out
// if the entire job has finished.
//
func (c *Coordinator) Done() bool {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	ret := true

	// Your code here.
	if len(c.MapTasks) > 0 && len(c.ReduceTasks) == 0 {
		ret = false
	} else {
		for _, task := range c.ReduceTasks {
			if !task.Finished {
				ret = false
				break
			}
		}
	}

	return ret
}

//
// create a Coordinator.
// main/mrcoordinator.go calls this function.
// nReduce is the number of reduce tasks to use.
//
func MakeCoordinator(files []string, nReduce int) *Coordinator {
	c := Coordinator{
		MapTasks:    make(map[int]*Task),
		ReduceTasks: make(map[int]*Task),
	}

	// Your code here.

	for i, file := range files {
		c.MapTasks[i] = &Task{
			Type:           MAP_TASK,
			MapId:          i,
			InputFilenames: []string{file},
			NReduce:        nReduce,
		}
	}

	for i := 0; i < nReduce; i++ {
		c.ReduceTasks[i] = &Task{
			Type:     REDUCE_TASK,
			ReduceId: i,
			NReduce:  nReduce,
		}
	}

	c.server()
	return &c
}
